/*
 * Description: defender exclusion handler
 */
#include "defender_exclusion_accelerator.h"
#include "Windows.h"
#include <tchar.h>
#include <algorithm>
#include <atlsafe.h>
#include <comutil.h>
#include <functional>
#include <vector>
#include "log.h"
#include "performance/utils/registry_utils.h"
#include "performance/utils/perf_utils.h"
#include "performance/perf_constants.h"
#include "performance/utils/defender_utils.h"

#pragma comment(lib, "wbemuuid.lib")
#pragma comment(lib, "comsuppw.lib")

namespace PERF {

bool DefenderExclusionAccelerator::OnBegin(std::shared_ptr<PerfInfo> perfInfo)
{
    LOG_INFO("[PERF - OnBegin] - start.");
    if (perfInfo == nullptr) {
        LOG_ERR("[PERF - OnBegin] - perfInfo is null!");
        return false;
    }

    bool ret = true;
    std::string tmpPath = perfInfo->appPath;
    std::transform(tmpPath.begin(), tmpPath.end(), tmpPath.begin(), tolower);

    std::vector<std::string> exclusionProcesses;
    DefenderUtils::GetCurrentExclusionItems(exclusionProcesses, ExclusionType::PROCESS);
    PERF_UTILS.ToLower(exclusionProcesses);
    if (std::count(exclusionProcesses.begin(), exclusionProcesses.end(), tmpPath) == 0) {
        ret &= DefenderUtils::AddExclusionItem(perfInfo->appPath, ExclusionType::PROCESS);
    } else {
        perfInfo->skipSign |= SKIP_SIGN_EXCLUSION_PROCESS;
    }

    std::vector<std::string> exclusionPaths;
    DefenderUtils::GetCurrentExclusionItems(exclusionPaths, ExclusionType::PATH);
    PERF_UTILS.ToLower(exclusionPaths);
    if (std::count(exclusionPaths.begin(), exclusionPaths.end(), tmpPath) == 0) {
        ret &= DefenderUtils::AddExclusionItem(perfInfo->appPath, ExclusionType::PATH);
    } else {
        perfInfo->skipSign |= SKIP_SIGN_EXCLUSION_PATH;
    }
    CoUninitialize();
    return ret;
}

bool DefenderExclusionAccelerator::OnEnd(std::shared_ptr<PerfInfo> perfInfo)
{
    LOG_INFO("[PERF - OnEnd] - start.");
    if (perfInfo == nullptr) {
        LOG_ERR("[PERF - OnEnd] - perfInfo is null!");
        return false;
    }
    bool ret = true;
    if (!(perfInfo->skipSign & SKIP_SIGN_EXCLUSION_PROCESS)) {
        ret &= DefenderUtils::RemoveExclusionItem(perfInfo->appPath, ExclusionType::PROCESS);
    }
    if (!(perfInfo->skipSign & SKIP_SIGN_EXCLUSION_PATH)) {
        ret &= DefenderUtils::RemoveExclusionItem(perfInfo->appPath, ExclusionType::PATH);
    }
    CoUninitialize();
    return ret;
}

}  // namespace PERF